import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import umap
from matplotlib import rcParams
from sklearn.decomposition import PCA
from sklearn.feature_selection import VarianceThreshold
from sklearn.model_selection import KFold, train_test_split
from sklearn.metrics import roc_auc_score
import torch
import torch.nn as nn
import torch.optim as optim

# 设置中文字体
rcParams['font.sans-serif'] = ['SimHei']  # 使用 SimHei 字体
rcParams['axes.unicode_minus'] = False   # 确保负号正常显示

# 数据加载
positive_data = pd.read_csv('data/positive.csv', header=None)
negative_data = pd.read_csv('data/negative.csv', header=None)

# 添加标签
positive_data['label'] = 1
negative_data['label'] = 0
data = pd.concat([positive_data, negative_data], ignore_index=True)

X = data.iloc[:, :-1].values
y = data.iloc[:, -1].values

# 1. 原始数据可视化 (UMAP 降维)
umap_model = umap.UMAP(n_neighbors=15, min_dist=0.1, metric='euclidean')
umap_embedding = umap_model.fit_transform(X)

plt.figure(figsize=(10, 8))
sns.scatterplot(x=umap_embedding[:, 0], y=umap_embedding[:, 1], hue=y, palette='viridis')
plt.title('UMAP Projection of Original Data')
plt.xlabel('UMAP 1')
plt.ylabel('UMAP 2')
plt.show()


# 定义BP神经网络
class BPNeuralNetwork(nn.Module):
    def __init__(self, input_dim):
        super(BPNeuralNetwork, self).__init__()
        self.layer1 = nn.Linear(input_dim, 128)
        self.layer2 = nn.Linear(128, 32)
        self.layer3 = nn.Linear(32, 2)
        self.output_layer = nn.Linear(2, 1)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.relu(self.layer1(x))
        x = self.relu(self.layer2(x))
        x = self.relu(self.layer3(x))
        x = self.sigmoid(self.output_layer(x))
        return x


# 手动划分数据为五折
positive_indices = np.where(y == 1)[0]
negative_indices = np.where(y == 0)[0]
np.random.shuffle(positive_indices)
np.random.shuffle(negative_indices)

folds = []
for i in range(5):
    pos_fold = positive_indices[i * 40:(i + 1) * 40]
    neg_fold = negative_indices[i * 40:(i + 1) * 40]
    folds.append(np.concatenate([pos_fold, neg_fold]))

# 2. 五折交叉验证
X_tensor = torch.tensor(X, dtype=torch.float32)
y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1)
auroc_scores = []

for i in range(5):
    test_indices = folds[i]
    train_indices = np.concatenate([folds[j] for j in range(5) if j != i])

    X_train, X_test = X_tensor[train_indices], X_tensor[test_indices]
    y_train, y_test = y_tensor[train_indices], y_tensor[test_indices]

    model = BPNeuralNetwork(X.shape[1])
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(100):
        model.train()
        optimizer.zero_grad()
        output = model(X_train)
        loss = criterion(output, y_train)
        loss.backward()
        optimizer.step()

    model.eval()
    with torch.no_grad():
        y_pred = model(X_test).squeeze().numpy()
        auc = roc_auc_score(y_test.numpy(), y_pred)
        auroc_scores.append(auc)
        print(f'Fold {i + 1}: AUROC = {auc:.4f}')

        # 导出倒数第二层数据
        intermediate_output = model.relu(model.layer2(model.relu(model.layer1(X_test)))).numpy()

        # 绘制测试集散点图
        plt.figure(figsize=(14, 6))

        # 子图1：预测标签
        plt.subplot(1, 2, 1)
        predicted_labels = (y_pred > 0.5).astype(int)
        plt.scatter(intermediate_output[:, 0], intermediate_output[:, 1], c=predicted_labels, cmap='viridis')
        plt.title(f'Fold {i + 1}: 预测标签')
        plt.xlabel('Dim 1')
        plt.ylabel('Dim 2')

        # 子图2：真实标签
        plt.subplot(1, 2, 2)
        plt.scatter(intermediate_output[:, 0], intermediate_output[:, 1], c=y_test.numpy(), cmap='viridis')
        plt.title(f'Fold {i + 1}: 真实标签')
        plt.xlabel('Dim 1')
        plt.ylabel('Dim 2')

        plt.show()

print(f'平均 AUROC: {np.mean(auroc_scores):.4f}')

#3. 打乱数据重新划分训练集和测试集
np.random.seed(42)
shuffled_indices = np.random.permutation(len(X))
X_shuffled = X[shuffled_indices]
y_shuffled = y[shuffled_indices]

X_train, X_test, y_train, y_test = train_test_split(X_shuffled, y_shuffled, test_size=0.2, random_state=42)

X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32).view(-1, 1)

# 重新训练模型
model = BPNeuralNetwork(X.shape[1])
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    output = model(X_train_tensor)
    loss = criterion(output, y_train_tensor)
    loss.backward()
    optimizer.step()

# 测试模型性能
model.eval()
with torch.no_grad():
    y_pred = model(X_test_tensor).squeeze().numpy()
    auc = roc_auc_score(y_test_tensor.numpy(), y_pred)
    print(f'打乱正负样本排列后重新划分训练模型: AUROC = {auc:.4f}')


# 绘制 PCA 后的散点图

# 3. 特征选择（方差最大的两个特征）
selector = VarianceThreshold()
variances = selector.fit(X).variances_
top_two_indices = np.argsort(variances)[-2:]
X_selected = X[:, top_two_indices]

plt.figure(figsize=(10, 8))
sns.scatterplot(x=X_selected[:, 0], y=X_selected[:, 1], hue=y, palette='viridis')
plt.title('经过特征选择后的散点图')
plt.xlabel('Feature 1')
plt.ylabel('Feature 2')
plt.show()

# 4. 特征提取（PCA 降维到二维）
pca = PCA(n_components=2)
X_pca = pca.fit_transform(X)

plt.figure(figsize=(10, 8))
sns.scatterplot(x=X_pca[:, 0], y=X_pca[:, 1], hue=y, palette='viridis')
plt.title('经过PCA降维后的散点图')
plt.xlabel('Principal Component 1')
plt.ylabel('Principal Component 2')
plt.show()